import os
from itertools import product

# Fixed parameters
CTA_Q = 128
CTA_K = 64
WARP_K = 64
DTypePVAccum = "float"
use_inst_buffer = True

# Swept parameters
head_dims = [64, 128]
qk_quant_grans = [1, 2, 3]
pv_threshold_modes = [0, 1]
dtypes_out = ["half", "nv_bfloat16"]
is_causals = [True, False]
return_pv_counts = [True, False]

# Output directory
output_dir = os.path.dirname(os.path.abspath(__file__))
script_name = os.path.basename(__file__)
header = f'// this file is automatically generated by running `python {script_name}`\n#include "../qk_int_sv_f16_cuda_sm80.cuh"\n'

def bool_to_int(b):
    return "1" if b else "0"

# Function parameter list
param_list = (
    "  int8_t* Q, int8_t* K, half* V, {dtype_out}* O,\n"
    "  int32_t* PV_Count, int32_t *__restrict__ Lut, int32_t *__restrict__ Valid_Block_Num, float *__restrict__ PV_Threshold,\n"
    "  float* Q_scale, float* K_scale,\n"
    "  const uint32_t batch_size, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_qo_heads, const uint32_t num_kv_heads,\n"
    "  const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q,\n"
    "  const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k,\n"
    "  const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v,\n"
    "  const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o,\n"
    "  float sm_scale\n"
)

# Generate combinations
for hd, qkg, pv_mode, dtype_out, causal, ret_pv_count in product(
        head_dims, qk_quant_grans, pv_threshold_modes, dtypes_out, is_causals, return_pv_counts):
    if ret_pv_count and pv_mode == 0:
        continue
    WARP_Q = 32 if hd == 64 else 16
    filename = (
        f"inst_sm80_ctaq{CTA_Q}_ctak{CTA_K}_warpq{WARP_Q}_warpk{WARP_K}"
        f"_hd{hd}_qkg{qkg}_pvacc{DTypePVAccum}_ibuf{bool_to_int(use_inst_buffer)}"
        f"_pvth{pv_mode}_o{dtype_out}_causal{bool_to_int(causal)}_retpvth{bool_to_int(ret_pv_count)}.cu"
    )
    filepath = os.path.join(output_dir, filename)

    instantiation = (
        f"template void SpargeAttentionSM80Dispatched<"
        f"{CTA_Q}, {CTA_K}, {WARP_Q}, {WARP_K}, {hd}, {qkg}, {DTypePVAccum}, "
        f"{str(use_inst_buffer).lower()}, {pv_mode}, {dtype_out}, "
        f"{str(causal).lower()}, {str(ret_pv_count).lower()}"
        f">(\n{param_list.format(dtype_out=dtype_out)});"
    )

    with open(filepath, "w") as f:
        f.write(header)
        f.write(instantiation + "\n")

print(f"Generated {len(os.listdir(output_dir))} instantiations in '{output_dir}'")
